///
/// Copyright (C) 2016, Dependable Systems Laboratory, EPFL
/// Copyright (C) 2016, Cyberhaven
///
/// Permission is hereby granted, free of charge, to any person obtaining a copy
/// of this software and associated documentation files (the "Software"), to deal
/// in the Software without restriction, including without limitation the rights
/// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
/// copies of the Software, and to permit persons to whom the Software is
/// furnished to do so, subject to the following conditions:
///
/// The above copyright notice and this permission notice shall be included in all
/// copies or substantial portions of the Software.
///
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
/// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
/// SOFTWARE.
///

#ifndef S2E_PLUGINS_Recipe_H_
#define S2E_PLUGINS_Recipe_H_

#include <s2e/cpu.h>

#include <klee/Expr.h>
#include <s2e/Plugin.h>
#include <s2e/Plugins/ExecutionMonitors/StackMonitor.h>
#include <s2e/Plugins/OSMonitors/OSMonitor.h>
#include <s2e/Plugins/OSMonitors/Support/MemUtils.h>
#include <s2e/Plugins/OSMonitors/Support/MemoryMap.h>
#include <s2e/Plugins/OSMonitors/Support/ModuleMap.h>
#include <s2e/Plugins/OSMonitors/Support/ProcessExecutionDetector.h>
#include <s2e/Plugins/Searchers/SeedSearcher.h>
#include <s2e/Plugins/Support/KeyValueStore.h>

#include "RecipeDescriptor.h"

namespace s2e {
namespace plugins {
namespace recipe {

struct MemPrecondition {
    Preconditions preconditions;
    klee::ref<klee::Expr> ptrExpr;
    unsigned requiredMemSize;
    bool exec; // must be executable
    MemPrecondition() : requiredMemSize(0), exec(false) {
    }
};

typedef std::map<std::string, RecipeDescriptor *> RecipeMap;

struct RecipeStats {
    /// Number of recipes that couldn't be parsed
    unsigned invalidRecipeCount;

    /// Number of times all recipes failed
    /// to apply to an execution state
    unsigned failedRecipeTries;

    /// Number of times at least one recipe
    /// could be applied to an execution state
    unsigned successfulRecipeTries;

    RecipeStats() {
        invalidRecipeCount = 0;
        failedRecipeTries = 0;
        successfulRecipeTries = 0;
    }
};

inline llvm::raw_ostream &operator<<(llvm::raw_ostream &out, const RecipeStats &s) {
    out << "RecipeStats invalidCount: " << s.invalidRecipeCount << " failedTries: " << s.failedRecipeTries
        << " successTries: " << s.successfulRecipeTries;

    return out;
}

class Recipe : public Plugin {
    S2E_PLUGIN
public:
    Recipe(S2E *s2e) : Plugin(s2e) {
    }

    void initialize();

    sigc::signal<void, S2EExecutionState *, const PovOptions &, const std::string & /* recipeName */> onPovReady;

private:
    using time_point = std::chrono::steady_clock::time_point;

    OSMonitor *m_monitor;
    ProcessExecutionDetector *m_process;
    ModuleMap *m_modules;
    StackMonitor *m_stackMonitor;
    seeds::SeedSearcher *m_seedSearcher;
    KeyValueStore *m_keyValueStore;
    MemoryMap *m_map;
    MemUtils *m_memutils;

    RecipeMap m_recipes;
    time_point m_lastRecipeLoadTime;
    std::string m_recipesDir;

    RecipeStats m_stats;

    uint32_t m_flagPage;
    bool m_abortOnInvalidRecipe;

    void injectSegFault(S2EExecutionState *state);

    void onTimer();

    void loadRecipesFromDirectory(const std::string &directory);

    void instrumentCTI(ExecutionSignal *signal, S2EExecutionState *state, TranslationBlock *tb, uint64_t pc,
                       bool isStatic, uint64_t staticTarget);

    void onTranslateBlockEnd(ExecutionSignal *signal, S2EExecutionState *state, TranslationBlock *tb, uint64_t pc,
                             bool isStatic, uint64_t staticTarget);

    void onAfterCall(S2EExecutionState *state, uint64_t callInstructionPc);

    void handleICTI(S2EExecutionState *state, uint64_t pc, unsigned rm, int op, int offset);
    void onTranslateICTIStart(ExecutionSignal *signal, S2EExecutionState *state, TranslationBlock *tb, uint64_t pc,
                              int rm, int op, int offset);

    void onSymbolicAddress(S2EExecutionState *state, klee::ref<klee::Expr> virtualAddress, uint64_t concreteAddress,
                           bool &concretize, CorePlugin::symbolicAddressReason reason);

    void onBeforeSymbolicDataMemoryAccess(S2EExecutionState *state, klee::ref<klee::Expr> addr,
                                          klee::ref<klee::Expr> value, bool isWrite);

    bool exprIn(klee::ref<klee::Expr> expr, const ExprList &a);
    void extractSymbolicBytes(const klee::ref<klee::Expr> &e, ExprList &bytes);

    klee::ref<klee::Expr> getRegbyteExpr(S2EExecutionState *state, const StateConditions &sc,
                                         const klee::ref<Register> &r);
    klee::ref<klee::Expr> getRegExpr(S2EExecutionState *state, const StateConditions &sc, const klee::ref<Register> &r);
    bool getLeftExpr(S2EExecutionState *state, const StateConditions &sc, const Precondition &p,
                     klee::ref<klee::Expr> &left);
    bool isSymbolicRegPtr(S2EExecutionState *state, const StateConditions &sc, const klee::ref<Left> &l,
                          klee::ref<klee::Expr> &ptrExpr);

    void pruneSymbolicSequences(S2EExecutionState *state, const ExprList &usedExprs,
                                std::vector<MemUtils::AddrSize> &list);

    bool applySimplePrecondition(S2EExecutionState *state, const StateConditions &sc, const klee::ref<klee::Expr> &left,
                                 const klee::ref<Right> &right, RecipeConditions &recipeConditions);
    bool testMemPrecondition(S2EExecutionState *state, const StateConditions &sc, const MemPrecondition &p,
                             MemUtils::AddrSize sequence, const RecipeConditions &recipeConditions, uint32_t &offset);
    bool applyMemPrecondition(S2EExecutionState *state, const StateConditions &sc, const MemPrecondition &p,
                              RecipeConditions &recipeConditions);

    void classifyPreconditions(S2EExecutionState *state, const StateConditions &sc, const Preconditions &p,
                               Preconditions &simple, std::map<Register::Reg, MemPrecondition> &memory);
    bool applyPreconditions(S2EExecutionState *state, PovType type, const StateConditions &sc, const Preconditions &p,
                            RecipeConditions &recipeConditions);
    bool tryRecipes(S2EExecutionState *state, const StateConditions &sc, RecipeConditions &recipeConditions);

    void tryRecipesOnICTI(S2EExecutionState *state, const klee::ref<klee::Expr> &regExpr,
                          const klee::ref<Register> &reg, int ictiOffset);

    bool checkUsedRegs(S2EExecutionState *state, const klee::ref<Left> &left, const RegList &usedRegs);

    bool getCurrentModule(S2EExecutionState *state, uint64_t eip, ModuleDescriptorConstPtr *module);

    void suppressExecutionWithInvalidAddress(S2EExecutionState *state, klee::ref<klee::Expr> addr, bool isWrite,
                                             int accessSize);
    void handleSymbolicWrite(S2EExecutionState *state, const klee::ref<klee::Expr> &addr);

    void handleSymbolicRead(S2EExecutionState *state, const klee::ref<klee::Expr> &addr);

public:
    const RecipeStats &getStats() const {
        return m_stats;
    }

    void resetStats() {
        m_stats = RecipeStats();
    }

    unsigned getRecipeCount() const {
        return m_recipes.size();
    }
};

} // namespace recipe
} // namespace plugins
} // namespace s2e

#endif /* S2E_PLUGINS_Recipe_H_ */
